from langchain.chains.llm import LLMChain
from langchain.chains.sequential import SimpleSequentialChain
from langchain.chains.transform import TransformChain
from langchain_community.llms.openai import OpenAI
from langchain_core.prompts import PromptTemplate
import __init__

with open("../kecheng源码/letter.txt") as f:
    letters = f.read()

def transform_func(inputs:dict) -> dict:
    text = inputs["text"]
    shortened_text = "\n\n".join(text.split("\n\n")[:3])
    return {"output_text": shortened_text}

# 文档转换链
transform_chain = TransformChain(
    input_variables=["text"],
    output_variables=["output_text"],
    transform=transform_func
)

template = """对下面的文字进行总结：
{output_text}

总结："""
prompt = PromptTemplate(
    input_variables=["output_text"],
    template=template
)
llm_chain = LLMChain(
    llm = OpenAI(),
    prompt = prompt
)
# 使用顺序链连接起来
sequential_chain=SimpleSequentialChain(
    chains=[transform_chain, llm_chain],
    verbose=True
)
print(letters)
print("--------------------------------------------------")
print(sequential_chain.run(letters))
